# classify_gui.py

import sys

import torch
import torch.nn as nn
from torchvision import models
import numpy as np
from PyQt5.QtWidgets import (
    QApplication, QWidget, QLabel, QPushButton, QFileDialog, QVBoxLayout, QMessageBox
)
from PyQt5.QtGui import QPixmap
from torchvision import transforms
from PIL import Image
import os
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


class ImageClassifier(QWidget):
    def __init__(self, model_path, idx_to_lable_path):
        super().__init__()
        self.setWindowTitle("图片分类器") # 设置窗口标题
        self.resize(400, 500) # 设置窗口大小

        self.model_path = model_path  # 模型路径处理

        self.idx_to_lable_path = idx_to_lable_path # 标签路径处理

        # 图像预处理
        self.transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor()
        ])

        # UI 元素
        self.image_label = QLabel("请加载一张图片")
        self.image_label.setFixedSize(300, 300)
        self.image_label.setStyleSheet("border: 1px solid black")
        self.image_label.setScaledContents(True)

        self.result_label = QLabel("预测结果：")
        self.load_button = QPushButton("加载图片")
        self.load_button.clicked.connect(self.load_image)

        # 布局
        layout = QVBoxLayout()
        layout.addWidget(self.image_label)
        layout.addWidget(self.load_button)
        layout.addWidget(self.result_label)
        self.setLayout(layout)

    def load_mode(self):
        # 加载模型
        # 载入模型结构
        self.model = models.resnet18(pretrained=True) # 载入预训练模型
        self.model.fc = nn.Linear(self.model.fc.in_features, 81) # 81类
        self.model.fc
        self.model  = self.model.to(device)


        # 加载模型权重
        state_dict = torch.load(self.model_path) 
        self.model.load_state_dict(state_dict)
        self.model.eval()
         


    def load_image(self):
        # 选择图片
        file_path, _ = QFileDialog.getOpenFileName(self, "选择图片", "", "Images (*.png *.jpg *.jpeg *.bmp)")
        # 显示图片
        if file_path:
            self.image_label.setPixmap(QPixmap(file_path))
            self.classify_image(file_path)

    def idx_to_lable(self):
        idx_to_labels = np.load(self.idx_to_lable_path, allow_pickle=True).item()
        return idx_to_labels

    def classify_image(self, image_path):
        try:
            self.load_mode() # 加载模型
            image = Image.open(image_path).convert("RGB") 
            image_tensor = self.transform(image).unsqueeze(0).to(device)
            with torch.no_grad():
                output = self.model(image_tensor)
                _, predicted = torch.max(output, 1)  # 修改：正确解包 torch.max 的返回值

            classes = self.idx_to_lable()
            '''  
                模型在GPU上运行，因此需要将 predicted.item() 转换为 CPU 上的整数。
            '''
            class_name = classes[predicted.detach().cpu().item()] 

            self.result_label.setText(f"预测结果：类别 {class_name}")
            
        except Exception as e:
            QMessageBox.critical(self, "错误", f"分类失败：{str(e)}")

if __name__ == "__main__":
    model_path = r'checkpoint_1\best-0.897.pth'
    idx_to_lable = r'idx_to_labels.npy'
    app = QApplication(sys.argv) # 创建应用程序对象
    window = ImageClassifier(model_path, idx_to_lable)
    window.show()
    sys.exit(app.exec_())